import numpy as np
import nrrd
import matplotlib.pyplot as plt
import mycode.Tool as Tool


class DataPreProcess:
    def __init__(self, path):
        self.INPUT_FOLDER = path
        # 建立一个列表，用于存储所有需要的nrrd文件的地址
        self.lstFilesNRRD = Tool.Load(path, '.nrrd').load_scan()

    # 读取并输出单张NRRD图像（分别从三个不同的截面）
    def read(self):
        # 读取第14个NRRD文件，准备输出最中间的一张
        readdata, header = nrrd.read(self.lstFilesNRRD[13])
        print('输出第八个文件的数据头信息：', header)
        print('输出第该文件每张图片的尺寸信息：', readdata.shape)
        image = np.asarray(readdata)
        # 横截面
        plt.imshow(image[:, :, int(image.shape[2] / 2)], cmap='bone')
        plt.show()
        # 冠状面
        plt.imshow(image[:, int(image.shape[1] / 2), :], cmap='autumn')
        plt.show()
        # 矢状面
        plt.imshow(image[int(image.shape[0] / 2), :, :], cmap='plasma')
        plt.show()

    def write(self):
        H = Tool.H5pyStore('heart_8', 'w')
        readdata, header = nrrd.read(self.lstFilesNRRD[8])
        H.createDataSet('image', np.asarray(readdata))
        H.close()

    def readFromHDF5(self):
        H = Tool.H5pyStore('heart_8', 'r')
        readdata = None
        for key in H.keys():
            readdata = H.get(key)
        print('从HDF5读出的数据：', readdata)
        print('数据具体内容：', readdata[()])
        H.close()

    def run(self):
        self.read()
        self.write()
        self.readFromHDF5()
